#include <cstdlib>
#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <sstream>

#ifdef COMPRESSION
#include <boost/iostreams/filter/zlib.hpp>
#include <boost/iostreams/filtering_stream.hpp>
namespace io = boost::iostreams;
#endif

#include <windows.h>

#include "path.h"
#include "archive.h"
#include "io.h"
#include "system.h"
#include "config.h"

namespace
{
    std::string g_exe_path;
    std::string g_launcher_path;
    std::string g_options;

    #ifdef COMPRESSION
    bool g_compress = false;
    #endif

    void print_usage_and_exit(const char *app_name)
    {
        #ifdef COMPRESSION
        #define COMPRESSION_SPEC " [/C]"
        #else
        #define COMPRESSION_SPEC ""
        #endif

        std::cout <<
            "EXE Packager version " VERSION "\n\n" <<
            app_name << " exefile" COMPRESSION_SPEC " [/L launcher] [/O options]\n\n"

            #ifdef COMPRESSION
            "  /C        Compress the packaged files. This usually results in a smaller     \n"
            "            output file, but your program will take longer to start.           \n"
            "                                                                               \n"
            #endif

            "  exefile   The EXE file to package. All files and folders in its folder will  \n"
            "            be packaged too.                                                   \n"
            "                                                                               \n"
            "  launcher  Where to write the launcher to.                                    \n"
            "                                                                               \n"
            "  options   A string of options to pass to the EXE file when it is executed.   \n"
            "            Additionally, the launcher will forward any options given to it to \n"
            "            the EXE file.                                                      \n"
            "                                                                               \n"
            "            The following variables are replaced when present in the options   \n"
            "            string:                                                            \n"
            "                    %W   The launcher's working directory                      \n"
            "                    %%   Literal percent sign                                  \n"
        ;
        std::exit(1);

        #undef COMPRESSION_SPEC
    }

    void fatal(const std::string &message)
    {
        std::cerr << message << "\n";
        std::exit(EXIT_FAILURE);
    }

    void parse_args(int argc, char **argv)
    {
        for(int i = 1; i < argc; i++)
        {
            char *arg = argv[i];

            if(!stricmp(arg, "/?"))
                print_usage_and_exit(argv[0]);
            else if(i == 1)
                g_exe_path = arg;
            else if(!stricmp(arg, "/L"))
                g_launcher_path = argv[++i];
            #ifdef COMPRESSION
            else if(!stricmp(arg, "/C"))
                g_compress = true;
            #endif
            else if(!stricmp(arg, "/O"))
            {
                for(int j=i+1; j<argc; j++)
                {
                    g_options.append(argv[j]);
                    if(j + 1 != argc)
                        g_options.append(" ");
                }
                break;
            }
        }
    }

    struct IconInfo
    {
        HMODULE exe;
        HANDLE launcher;
    };

    BOOL CALLBACK found_icon(HMODULE, LPCTSTR type, LPTSTR name, LONG_PTR p)
    {
        IconInfo &info  = *reinterpret_cast<IconInfo *>(p);
        HMODULE exe     = info.exe;
        HANDLE launcher = info.launcher;

        HRSRC res      = FindResource(exe, name, type);
        DWORD size     = SizeofResource(exe, res);
        HGLOBAL icon   = LoadResource(exe, res);
        void *data     = LockResource(icon);

        UpdateResource(launcher, type, name, 0, data, size);

        int e; if(e = GetLastError()) report_error_and_exit(e);

        return TRUE;
    }

    void insert_icons(HANDLE launcher)
    {
        LPTSTR       icon_type = MAKEINTRESOURCE(RT_ICON);
        LPTSTR group_icon_type = MAKEINTRESOURCE(RT_GROUP_ICON);

        IconInfo info;
        info.exe      = LoadLibrary(g_exe_path.c_str());
        info.launcher = launcher;

        LONG_PTR info_ptr = reinterpret_cast<LONG_PTR>(&info);

        EnumResourceNames(info.exe,       icon_type, found_icon, info_ptr);
        EnumResourceNames(info.exe, group_icon_type, found_icon, info_ptr);

        FreeLibrary(info.exe);
    }

    void write_empty_launcher()
    {
    	DWORD size;
		void *launcher_data = load_rcdata(size);
		std::ofstream launcher(g_launcher_path.c_str(), std::ios::binary);
		if(!launcher)
			fatal("Can't create launcher");
		launcher.write((char *)launcher_data, size);
    }

    void write_launcher()
    {
        std::string exe_dir = dir_name(g_exe_path);

        write_empty_launcher();

        std::string command =
            relative_path(exe_dir, g_exe_path) + " " + g_options;

		// Find all files to pack in
        Strings file_paths;
        find_files(exe_dir, file_paths);
        for(int i=0; i<file_paths.size(); i++)
            file_paths[i] = relative_path(exe_dir, file_paths[i]);

        // Create archive
        std::stringstream ss;
        #ifdef COMPRESSION
            io::filtering_stream<io::output> write;
            if (g_compress)
            {
                ss.put(1);
                write.push(io::zlib_compressor());
            }
            else
                ss.put(0);
            write.push(ss);
        #else
            std::stringstream &write = ss;
        #endif

        write_archive(write, exe_dir, file_paths, command);
        std::string archive = ss.str();
        unsigned size = archive.size();
        void *data = const_cast<char *>(archive.c_str());

        // Inject archive into launcher as a resource
        LPTSTR type = MAKEINTRESOURCE(RT_RCDATA);
        LPSTR name = MAKEINTRESOURCE(1);
        HANDLE launcher = BeginUpdateResource(g_launcher_path.c_str(), FALSE);
        UpdateResource(launcher, type, name, 0, data, size);
        insert_icons(launcher);
        EndUpdateResource(launcher, FALSE);
    }
}

int main(int argc, char **argv)
{
    parse_args(argc, argv);

    if(g_exe_path.empty())
    {
        print_usage_and_exit(argv[0]);
    }

    if(!file_exists(g_exe_path))
    {
        std::cerr << g_exe_path << " does not exist.\n";
        return 1;
    }

    if(g_launcher_path.empty())
    {
        g_launcher_path = bare_name(g_exe_path) + ".exe";
    }

    write_launcher();
}


